Replace op.Constant float scalars with Python literals for dtype-safe auto-casting by justinchuby · Pull Request #58 · onnxruntime/mobius

justinchuby · 2026-03-28T14:44:11Z

Problem

op.Constant(value_float=1.0) always creates a float32 ONNX Constant node. When the surrounding computation uses bfloat16 or float16 tensors, this introduces a dtype mismatch — ONNX requires both operands to have the same type.

Solution

onnxscript auto-casts Python scalars (int, float, bool) to match the dtype of the other operand in binary ops. So op.Add(bf16_tensor, 1.0) works correctly — the 1.0 is auto-cast to bf16.

This PR replaces op.Constant(value_float=...) with plain Python literals wherever possible. In two places where auto-cast cannot help (neg_inf values passed to op.Expand which has no typed operand to infer from), op.CastLike(-1e30, scores) is used instead.

Changes

File	Before	After
`_audio.py`	`op.Constant(value_float=0.5)`	`0.5` (Python literal)
`_diffusion.py`	`op.Constant(value_float=1.0)` → `one` variable	`1.0` inline
`_moe.py` (SigmoidTopKGate)	`op.CastLike(op.Constant(...), ...)` for eps, scale	`1e-9`, `self.routed_scaling_factor` (Python literals)
`_moe.py` (SparseMixerGate)	`op.Constant(value_float=...)` for threshold, neg_inf	`2.0 * jitter_eps` (literal); `op.CastLike(-1e30, scores)` for neg_inf
`gemma3n.py` (router)	`op.Constant(value_float=self.router_input_scale)`	`self.router_input_scale` (Python float)
`gemma3n.py` (Laurel residual)	`op.Constant(value_float=float(math.sqrt(2)))`	`float(math.sqrt(2))` (Python literal)

Testing

New test file: src/mobius/components/_castlike_dtype_test.py

7 parametrized tests across OffsetRMSNorm, quick_gelu, and AdaLayerNormOutput
Tests with FP16 and BF16 inputs verify: zero CastLike ops emitted (auto-cast handles everything) and output dtype matches input dtype
Float32 control test confirms no regression

Note

bfloat16 end-to-end inference still fails on ORT CPU EP (no bf16 kernels) and on ORT CUDA EP (Attention op decomposition creates mixed-type nodes internally) — these are ORT runtime limitations, not graph construction issues.

Copilot

Pull request overview

Fixes ONNX mixed-dtype load errors when adding a scalar 1.0 to tensors that may become BFLOAT16 after weight application (notably in OffsetRMSNorm and Gemma3n AltUp coefficient correction), by explicitly casting the scalar Constant to match the target tensor dtype.

Changes:

Cast the scalar 1.0 Constant via op.CastLike(..., weight_tensor) before op.Add in OffsetRMSNorm.
Apply the same cast-before-add pattern to Gemma3n AltUp correction coefficients.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
src/mobius/components/_rms_norm.py	Prevents dtype mismatch in `OffsetRMSNorm` by casting the scalar constant to the weight’s dtype before Add.
src/mobius/models/gemma3n.py	Prevents dtype mismatch in Gemma3n AltUp coefficient correction by casting the scalar constant to the coefficients’ dtype before Add.

github-actions · 2026-03-28T14:47:34Z

🏗️ Architecture Diff

Comparing 7c1972d → 228d01e

Model	Sub-model	Status
bert (feature-extraction)	model	⚪
falcon	model	⚪
gemma2	model	⚪
gpt2	model	⚪
llama	model	⚪
llama (static-cache)	model	⚪
mamba (ssm-text-generation)	model	⚪
phi3	model	⚪
phi3 (static-cache)	model	⚪
qwen	model	⚪
qwen (static-cache)	model	⚪
qwen2	model	⚪
qwen2 (static-cache)	model	⚪
qwen2_moe	model	⚪
qwen2_moe (static-cache)	model	⚪
qwen3	model	⚪
qwen3 (static-cache)	model	⚪
qwen3_5_moe (hybrid-text-generation)	model	⚪
qwen3_5_text (hybrid-text-generation)	model	⚪
qwen3_5_vl (hybrid-qwen-vl)	decoder	⚪
qwen3_5_vl (hybrid-qwen-vl)	embedding	⚪
qwen3_5_vl (hybrid-qwen-vl)	vision	⚪
qwen3_moe	model	⚪
qwen3_moe (static-cache)	model	⚪
qwen3_next (hybrid-text-generation)	model	⚪
t5 (seq2seq)	decoder	⚪
t5 (seq2seq)	encoder	⚪
whisper (speech-to-text)	decoder	⚪
whisper (speech-to-text)	encoder	⚪

No architecture changes detected. ✅

Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

Use CastLike to cast the float constant 1.0 to match the weight/coefs dtype before Add. Without this, op.Add(bf16_tensor, 1.0) creates a type mismatch since 1.0 becomes a float32 Constant. Note: bf16 inference still fails on ORT CPU EP (no bf16 kernels) and on CUDA EP due to ORT Attention op decomposition creating mixed-type nodes internally. This is an ORT limitation, not a graph construction issue. f16 works on both CPU and CUDA. Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Ensure float32 literals are cast to match model-precision tensors before use in arithmetic ops. Affected sites: - _activations.py: quick_gelu scalar 1.702 - _diffusion.py: AdaLayerNormZero constant 1.0 added to scale - _moe.py: threshold and neg_inf constants matched to scores dtype (x2) - gemma3n.py: router_input_scale constant; sqrt(2) divisor in Laurel residual - _ssm.py: epsilon constant in RMS variance computation - _audio.py: 0.5 half-scaling constant in ConformerBlock Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Tests build OffsetRMSNorm, quick_gelu, and AdaLayerNormOutput with float16/bfloat16 inputs and assert CastLike ops are present. Verified tests FAIL without the CastLike fix and PASS with it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

onnxscript auto-casts Python scalars (int/float/bool) to match the dtype of the other operand in binary ops. op.Constant(value_float=x) returns ir.Value(FLOAT) which is NOT auto-cast, causing dtype mismatches in BF16/FP16 builds. Changes: - _activations.py: 1.702 → Python literal (op.Mul auto-casts) - _audio.py: 0.5 → inline literals (removes CastLike + intermediate var) - _diffusion.py: 1.0 → Python literal (op.Add auto-casts) - _moe.py: 1e-9, routed_scaling_factor, 2*jitter_eps → Python literals; -1e30 uses op.CastLike(-1e30, scores) with Python literal to avoid cache-key collision (Expand has no type-variable binding for its input) - _rms_norm.py: 1.0 → Python literal (op.Add auto-casts) - gemma3n.py: router_input_scale, 1.0, sqrt(2) → Python literals Test update: TestCastLikeDtypeSafety → TestPythonLiteralAutocast. New tests call _cast_module_dtype to simulate production dtype promotion, then verify output dtype matches input dtype with no CastLike ops present. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

github-actions · 2026-04-01T17:44:51Z

Performance Comparison

Comparing 7c1972d → 228d01e

Model	Metric	Baseline	Current	Delta
bert (feature-extraction)	model_size_bytes	359 KB	359 KB	+0.0%	⚪
bert (feature-extraction)	num_nodes	61	61	+0.0%	⚪
falcon	model_size_bytes	364 KB	364 KB	+0.0%	⚪
falcon	num_nodes	66	66	+0.0%	⚪
gemma2	model_size_bytes	428 KB	428 KB	+0.0%	⚪
gemma2	num_nodes	107	107	+0.0%	⚪
gpt2	model_size_bytes	388 KB	388 KB	+0.0%	⚪
gpt2	num_nodes	53	53	+0.0%	⚪
llama	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama	num_nodes	61	61	+0.0%	⚪
llama (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama (static-cache)	num_nodes	58	58	+0.0%	⚪
mamba (ssm-text-generation)	model_size_bytes	360 KB	360 KB	+0.0%	⚪
mamba (ssm-text-generation)	num_nodes	99	99	+0.0%	⚪
phi3	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3	num_nodes	59	59	+0.0%	⚪
phi3 (static-cache)	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3 (static-cache)	num_nodes	56	56	+0.0%	⚪
qwen2	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2	num_nodes	61	61	+0.0%	⚪
qwen2 (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2 (static-cache)	num_nodes	58	58	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	model_size_bytes	506 KB	506 KB	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	num_nodes	275	275	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	model_size_bytes	458 KB	458 KB	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	num_nodes	129	129	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	model_size_bytes	977 KB	977 KB	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	num_nodes	409	409	+0.0%	⚪
t5 (seq2seq)	model_size_bytes	836 KB	836 KB	+0.0%	⚪
t5 (seq2seq)	num_nodes	174	174	+0.0%	⚪
whisper (speech-to-text)	model_size_bytes	1008 KB	1008 KB	+0.0%	⚪
whisper (speech-to-text)	num_nodes	128	128	+0.0%	⚪

No performance regressions.

Add 10 new parametrized tests covering the 5 components identified by the code reviewer as missing BF16/FP16 coverage: SigmoidTopKGate (2 tests x 2 dtypes = 4): - test_routing_weights_dtype: verifies 1e-9 epsilon in op.Add auto-casts so routing weights stay in FP16/BF16 (not widened to FP32) - test_routed_scaling_factor_autocasts: verifies routed_scaling_factor Python float literal in op.Mul auto-casts to routing dtype SparseMixerGate (1 test x 2 dtypes = 2): - test_castlike_neg_inf_uses_input_dtype: verifies op.CastLike(-1e30, scores) correctly casts the -1e30 constant to the input dtype for use in op.Where and op.Expand, preventing FP32 type mismatches ConformerEncoderLayer (1 test x 2 dtypes = 2): - test_macaron_half_weight_autocasts: verifies 0.5 Macaron weight literals in op.Mul auto-cast to hidden state dtype (FP16/BF16) Gemma3nAltUp (1 test x 2 dtypes = 2): - test_router_input_scale_autocasts: verifies router_input_scale (hidden_size**-1.0 Python float) auto-casts in op.Mul during _compute_router_modalities All 17 tests pass (10 new + 7 existing). Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

codecov · 2026-04-13T17:53:46Z

Codecov Report

❌ Patch coverage is 98.40000% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/mobius/components/_castlike_dtype_test.py	98.26%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

…rals Replace all 68 instances of op.Constant(value_float=X) across 30 files with plain Python float literals or float() expressions. This leverages onnxscript's auto-casting of Python scalars to match the dtype of the other operand in binary ops, which fixes dtype mismatches when models use bfloat16 or float16. Patterns replaced: - op.Constant(value_float=X) -> X (plain literal) - op.Constant(value_float=float(expr)) -> float(expr) - op.Constant(value_float=self.attr) -> self.attr - op.CastLike(op.Constant(value_float=X), ref) -> op.CastLike(X, ref) Note: op.Constant(value_int=X) replacements were NOT included because onnxscript creates initializers (not Constant nodes) for Python int literals, causing 'already registered' collisions when the same int value appears multiple times in a graph. This is an onnxscript limitation that needs to be resolved upstream before value_int cleanup can proceed. Continuation of PR #58 which established this pattern. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>

…autocast (#166) Continuation of PR #58 which established this pattern. Replace all 68 instances of op.Constant(value_float=X) across 30 files with plain Python float literals or float() expressions. This leverages onnxscript's auto-casting of Python scalars to match the dtype of the other operand in binary ops, which fixes dtype mismatches when models use bfloat16 or float16. Patterns replaced: - op.Constant(value_float=X) -> X (plain literal) - op.Constant(value_float=float(expr)) -> float(expr) - op.Constant(value_float=self.attr) -> self.attr - op.CastLike(op.Constant(value_float=X), ref) -> op.CastLike(X, ref) Note: This is part 1. (A similar change can be done for int/ints/floats, but will do that in separate PRs.) --------- Signed-off-by: G Ramalingam <grama@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

justinchuby requested review from a team and Copilot March 28, 2026 14:44

Copilot started reviewing on behalf of justinchuby March 28, 2026 14:44 View session

justinchuby added the ai Created by an AI agent label Mar 28, 2026

Copilot AI reviewed Mar 28, 2026

View reviewed changes

Comment thread src/mobius/components/_rms_norm.py Outdated

Comment thread src/mobius/models/gemma3n.py Outdated

justinchuby force-pushed the justinchu/fix-castlike-dtype branch from a931ea8 to 67a3842 Compare March 28, 2026 14:48

justinchuby marked this pull request as draft March 28, 2026 15:46

justinchuby and others added 2 commits April 1, 2026 09:23

justinchuby force-pushed the justinchu/fix-castlike-dtype branch from a440239 to 4d56eca Compare April 1, 2026 16:25

justinchuby and others added 2 commits April 1, 2026 09:46

justinchuby marked this pull request as ready for review April 1, 2026 18:42

justinchuby changed the title ~~Fix bfloat16/float type mismatch in OffsetRMSNorm and Gemma3nAltUp~~ Replace op.Constant float scalars with Python literals for dtype-safe auto-casting Apr 1, 2026

gramalingam approved these changes Apr 13, 2026

View reviewed changes

Merge branch 'main' into justinchu/fix-castlike-dtype

228d01e

gramalingam merged commit 0df5e3f into main Apr 13, 2026
20 of 22 checks passed

gramalingam deleted the justinchu/fix-castlike-dtype branch April 13, 2026 18:41

gramalingam mentioned this pull request Apr 13, 2026

Replace op.Constant(value_float=...) with Python float literals, use autocast #166

Merged

Navigation Menu

Uh oh!

Replace op.Constant float scalars with Python literals for dtype-safe auto-casting - #58

Replace op.Constant float scalars with Python literals for dtype-safe auto-casting#58
gramalingam merged 6 commits into
mainfrom
justinchu/fix-castlike-dtype

justinchuby commented Mar 28, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Mar 28, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Apr 1, 2026 •

edited

Loading

Uh oh!

codecov Bot commented Apr 13, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

justinchuby commented Mar 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Solution

Changes

Testing

Note

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Mar 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🏗️ Architecture Diff

Uh oh!

github-actions Bot commented Apr 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Performance Comparison

Uh oh!

codecov Bot commented Apr 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

justinchuby commented Mar 28, 2026 •

edited

Loading

github-actions Bot commented Mar 28, 2026 •

edited

Loading

github-actions Bot commented Apr 1, 2026 •

edited

Loading

codecov Bot commented Apr 13, 2026 •

edited

Loading